﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;
using System.Diagnostics;
using Drifting.MotionPlanning.Primitives;
using Drifting.MotionPlanning.Systems;
using Drifting.MotionPlanning.PointSets;

namespace Drifting.MotionPlanning.SingleThreading.Algorithms
{
    public class iSST : TreeBasedPlanningAlgorithm
    {
        protected PointSet Vsupport;
        protected PointSet Vactive;
        protected double dv;
        protected double ds;

        public iSST(BaseSystem system, Random rand, double dv, double ds)
            : base(system, rand)
        {
            int[] branches = new int[] { 8, 8, 8, 8, 8, 8 };   //{ 6, 6, 6, 6, 6, 6, 6, 6, 6 }
            Vactive = new BinQuery(system.Min_Alias, system.Max_Alias, system.DOF_Alias, branches);
            Vsupport = new BinQuery(system.Min_Alias, system.Max_Alias, system.DOF_Alias, branches);
            vertices = Vactive;

            this.dv = dv;
            this.ds = ds;

            SetRoot(system.GenerateZeroState());
        }

        override public void AliasOf(double[] q, ref double[] alias)
        {
            if (alias == null)
                alias = new double[system.DOF_Alias];
            for (int i = 0; i < system.DOF_Alias; ++i)
                alias[i] = q[i];
        }

        override public void Execute(object input)
        {
            if (input is Int32 == false)
                throw new ArgumentException();
            n_exec = (Int32)input;
            if (x_root == null)
                throw new InvalidOperationException("Forget to call SetRoot first?");

            double[] a_sample = null;
            n_current = 0;
            List<State> neighbors = null;
            List<State> X_near = new List<State>(128);

            int n_KNeighbor = 0;
            int n_BN = 0;
            double controllerDuration = 10;

            ///////////////////////////////////
            //                               //
            bool enableKTest = true;         //
            bool enableBestNear = true;      //
            bool enableDrain = true;         //
            //                               //
            ///////////////////////////////////
            clearIPS();

            while (n_current < n_exec)
            {
                n_current++;

                #region Break point
                if (n_current > 5000 && n_current % 100 == 0)
                {
                    calculateIPS();

                    List<State> train = new List<State>();
                    Vactive.ReportVerticesSnapshot(ref train);

                    double aveCost = 0;
                    for (int i = 0; i < train.Count; ++i)
                    {
                        aveCost += train[i].Cost;
                    }
                    aveCost /= train.Count;
                    Debugger.Log(1, "checkpoint", "ips = " + ips +", aveCost = " + aveCost + ", iteration = " + n_current + Environment.NewLine);
                }
                #endregion

                #region Sample
                //Sample
                system.SampleAlias(rand, ref a_sample);
                #endregion

                #region Range query

                State x_BestNear = null;
                double cost_min = double.MaxValue;
                if (enableBestNear)
                {

                    //Range query
                    Vactive.Neighbor_RangeQuery(a_sample, dv, ref neighbors);

                    //Find the lowest cost state in range query
                    for (int i = 0; i < neighbors.Count; ++i)
                    {
                        if (neighbors[i].Cost < cost_min)
                        {
                            cost_min = neighbors[i].Cost;
                            x_BestNear = neighbors[i];
                        }
                    }
                }

                if (x_BestNear == null)
                {
                    double dist2_sample = double.MaxValue;
                    Vactive.Neighbor_SingleQuery(a_sample, ref x_BestNear, ref dist2_sample);
                }
                #endregion

                #region Propagate from x_BestNear
                //Propagate from x_BestNear
                double[] iCtrl = null;
                double iDt = double.MaxValue;
                system.SampleControl(rand, ref iCtrl, ref iDt);

                double[] q_new = null;
                double iCost = double.MaxValue;
                bool propagate_success = system.Propagate(memoryPad, x_BestNear.Q, iCtrl, iDt, ref q_new, ref iCost);

                if (propagate_success == false)
                    continue;
                               
                State x_new = null;
                Edge e_new = null;
                x_new = SpawnState(q_new);
                e_new = SpawnEdge(x_BestNear, x_new, iCtrl, iDt, iCost);
                if (system.ValidateEdge(e_new) == false)
                    continue;

                x_new.ParentEdge = e_new;
                double cost_BN = x_BestNear.Cost + iCost;
                x_new.Cost = cost_BN;

                #endregion

                #region K query
                //K query
                int K = (int)((1 + (int)Math.Ceiling(Math.Log(vertices.Count))) * 10);
                vertices.Neighbor_KQuery(a_sample, K, ref X_near);
                double kRadius = 0;
                for (int i = 0; i < X_near.Count; ++i)
                {
                    State x_i = X_near[i];
                    double dist_i = Math.Sqrt(PointSet.Euclidean2(a_sample, x_i.Alias));
                    if (kRadius < dist_i)
                    {
                        kRadius = dist_i;
                    }
                }
                #endregion

                #region Controller K tests

                double cost_KNeighbors_min = double.MaxValue;
                Edge e_KNeighbors_min = null;

                if (enableKTest)
                {
                    if (kRadius >= dv)
                    {
                        //K query radius is larger
                        for (int i = 0; i < X_near.Count; ++i)
                        {
                            State x_i = X_near[i];
                            Edge e_i = system.ExecuteControlTo(q_new, dv, x_i, controllerDuration);
                            if (e_i == null)
                                continue;

                            if (system.ValidateEdge(e_i) == false)
                                continue;
                            
                            double cost_i = x_i.Cost + e_i.step_cost;
                            if (cost_i < cost_KNeighbors_min)
                            {
                                cost_KNeighbors_min = cost_i;
                                e_KNeighbors_min = e_i;
                            }
                        }
                    }
                }
                #endregion

                #region Grow tree
                //grow tree
                if (cost_KNeighbors_min < cost_BN)
                {
                    n_KNeighbor++;
                    e_new = e_KNeighbors_min;
                    x_new = e_new.exit;
                }
                else
                {
                    n_BN++;
                }
                Vactive.AddState(x_new);
                #endregion
                
                #region Drain
                if (enableDrain)
                {
                    //delta_s drain
                    Vactive.Neighbor_RangeQuery(a_sample, ds, ref neighbors);
                    cost_min = double.MaxValue;
                    State x_min = null;
                    for (int i = 0; i < neighbors.Count; ++i)
                    {
                        State x_i = neighbors[i];
                        if (x_i.Cost < cost_min)
                        {
                            cost_min = x_i.Cost;
                            x_min = x_i;
                        }
                    }
                    for (int i = 0; i < neighbors.Count; ++i)
                    {
                        State x_i = neighbors[i];
                        if (x_min == x_i)
                            continue;
                        Vactive.DelState(x_i);
                    }
                }
                #endregion
            }
        }

    }
}
